import torch
import torch.nn as nn
from ..backbone import build_backbone
from ..classifier import build_clsfier


class ModelLoader(nn.Module):
    def __init__(self,config) -> None:
        super().__init__()
        self.backbone = build_backbone(config.arch.backbone)
        # print(self.backbone.weight.size(0))
        # print(torch.channels_last(self.backbone))
        config.arch.classifier['last_channel'] = self.backbone.last_channel
        self.clsfier = build_clsfier(config.arch.classifier)

    def forward(self,x):
        x = self.backbone(x)
        x = self.clsfier(x)
        return x